/*
 * Copyright (c) 2022 HiSilicon (Shanghai) Technologies CO., LIMITED.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/* -------------------------------------------------------------------------- */
/* --- PUBLIC DEPENDANCIES -------------------------------------------------- */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <signal.h>
#include <pthread.h>
#include <sys/prctl.h>
#include <math.h>

#include "hi_common.h"
#include "hi_comm_sys.h"
#include "hi_comm_svp.h"
#include "sample_comm.h"
/* -------------------------------------------------------------------------- */
/* --- COMMON DEPENDANCIES -------------------------------------------------- */
extern "C" {
#include "audio_rtmp.h"

#include "sample_comm_svp.h"
#include "sample_comm_nnie.h"
#include "sample_comm_ive.h"

/* -------------------------------------------------------------------------- */
/* --- MPP DEPENDANCIES ----------------------------------------------------- */
#include "ive_img.h"
#include "vgs_img.h"

/* -------------------------------------------------------------------------- */
/* --- PRIVATE DEPENDANCIES ------------------------------------------------- */
#include "uart.h"
#include "trace.h"
#include "sample_lcd.h"
#include "lbph_tracking.h"
#include "svp_nnie_application.h"
#include "main_process.h"

#include "ffmpeg_rtmp.h"
}

#include <opencv2/core.hpp>
#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

#include <iostream>
/* -------------------------------------------------------------------------- */
/* --- PRIVATE VARIABLES ---------------------------------------------------- */
AiPlugLib g_workPlug = {0};

static SAMPLE_VI_CONFIG_S s_stViConfig;
static SAMPLE_VO_CONFIG_S s_pstVoConfig;
static SAMPLE_IVE_SWITCH_S s_stRfcnSwitch = { HI_FALSE, HI_FALSE, HI_FALSE };
static HI_BOOL s_bNnieStopSignal = HI_FALSE;
static HI_BOOL s_bRtmpStopSignal = HI_FALSE;
static pthread_t s_hAudioThread = 0;
static pthread_t s_hNnieThread = 0;
static pthread_t s_hLbphThread = 0;
static pthread_t s_hUartThread = 0;

static SAMPLE_IVE_KCF_S s_stIveKcfInfo;

int uartFd = 0;

/* -------------------------------------------------------------------------- */
/* --- PUBLIC VARIABLES ----------------------------------------------------- */

/* -------------------------------------------------------------------------- */
/* --- PUBLIC FUNCTIONS ----------------------------------------------------- */
#define SAMPLE_IVE_MUTEX_INIT_LOCK(mutex)       \
    do {                                        \
        (void)pthread_mutex_init(&(mutex), NULL); \
    } while (0)
#define SAMPLE_IVE_MUTEX_LOCK(mutex)      \
    do {                                  \
        (void)pthread_mutex_lock(&(mutex)); \
    } while (0)
#define SAMPLE_IVE_MUTEX_UNLOCK(mutex)      \
    do {                                    \
        (void)pthread_mutex_unlock(&(mutex)); \
    } while (0)
#define SAMPLE_IVE_MUTEX_DESTROY(mutex)      \
    do {                                     \
        (void)pthread_mutex_destroy(&(mutex)); \
    } while (0)

/* exit process */
static HI_VOID Pause(HI_VOID)
{
    printf("---------------press double-click Enter key to exit!---------------\n");
    (void)getchar();
    (void)getchar();
}

static HI_VOID SAMPLE_IVE_KcfInit(SAMPLE_IVE_KCF_S *pstIveKcfInfo)
{
    (HI_VOID)memset_s(pstIveKcfInfo, sizeof(SAMPLE_IVE_KCF_S), 0, sizeof(SAMPLE_IVE_KCF_S));
    pstIveKcfInfo->enKcfProcStat = KCF_PROC_END;
    pstIveKcfInfo->bUpdateDetect = HI_FALSE;
    SAMPLE_IVE_MUTEX_INIT_LOCK(pstIveKcfInfo->ModifyMutex);
}

static HI_VOID SAMPLE_IVE_KcfDeInit(SAMPLE_IVE_KCF_S *pstIveKcfInfo)
{
    pstIveKcfInfo->enKcfProcStat = KCF_PROC_END;
    pstIveKcfInfo->bUpdateDetect = HI_FALSE;
    SAMPLE_IVE_MUTEX_DESTROY(pstIveKcfInfo->ModifyMutex);
}

/* function : load model and run ai process */
HI_S32 Yolov5DectionAiProcess(VIDEO_FRAME_INFO_S srcfrm, VIDEO_FRAME_INFO_S dstfrm)
{
    HI_S32 s32Ret;

    /* performances variables */
    struct timeval tm;
    _meas_time_start(&tm);

    /* First time come in thread, load moal and init params */
    if (g_workPlug.yolo_model == 0) {
        s32Ret = Load_Yolov5_Lite_Model(&g_workPlug.yolo_model, pcYoloModelName);
        if (s32Ret < 0) {
            g_workPlug.yolo_model = 0;
            SAMPLE_PRT("Load yolo classify model err, s32Ret=%#x\n", s32Ret);
            return s32Ret;
        }
    }

    VIDEO_FRAME_INFO_S resizeFrm;
    s32Ret = MppFrmResize(&srcfrm, &resizeFrm, YOLO_MODEL_WIDTH, YOLO_MODEL_HEIGHT);
    if (s32Ret != HI_SUCCESS) {
        SAMPLE_PRT("MppFrmResize FAIL, s32Ret=%#x\n", s32Ret);
        return s32Ret;
    }
    SAMPLE_IVE_MUTEX_LOCK(s_stIveKcfInfo.ModifyMutex);
    s32Ret = Yolov5DetectResnetClassifyCal(g_workPlug.yolo_model, &resizeFrm, &dstfrm, &s_stIveKcfInfo);
    SAMPLE_IVE_MUTEX_UNLOCK(s_stIveKcfInfo.ModifyMutex);
    if (s32Ret != HI_SUCCESS) {
        SAMPLE_PRT("Yolo cal single img FAI, s32Ret=%#x\n", s32Ret);
        return s32Ret;
    }

    MppFrmDestroy(&resizeFrm);

    _meas_time_stop(tm, __FUNCTION__);

    return s32Ret;
}

/* function : ai process */
static HI_VOID* GetVpssChnFrameAiProcess(HI_VOID* pArgs)
{
    HI_S32 s32Ret;
    VIDEO_FRAME_INFO_S stBaseFrmInfo;
    VIDEO_FRAME_INFO_S stExtFrmInfo;
    const HI_S32 s32MilliSec = 20000; /* 20000ms timeout */
    const VO_LAYER voLayer = 0;
    const VO_CHN voChn = 0;
    const VENC_CHN vencChn = 0;
    const HI_S32 s32VpssGrp = 0;
    const HI_S32 as32VpssChn[] = {VPSS_CHN1, VPSS_CHN2};

    hi_unused(pArgs);

    while (HI_FALSE == s_bNnieStopSignal) {
        /* 
         * VPSS 通道1 图像 用于画框图 以及 LBPH的输入图像裁剪 
         * VPSS channel 1 image used for frame drawing and LBPH input image cropping
         */
        s32Ret = HI_MPI_VPSS_GetChnFrame(0, 2, &stExtFrmInfo, s32MilliSec);
        if (s32Ret != HI_SUCCESS) {
            SAMPLE_PRT("Error(%#x),HI_MPI_VPSS_GetChnFrame failed, VPSS_GRP(%d), VPSS_CHN(%d)!\n", s32Ret, s32VpssGrp,
                as32VpssChn[1]);
            continue;
        }

        /* 
         * VPSS 通道0 图像 缩放至 320 * 320 后用于YOLOV5-LITE目标识别
         * VPSS channel 0 image is scaled to 320 * 320 for YOLOV5-LITE target recognition
         */
        s32Ret = HI_MPI_VPSS_GetChnFrame(0, 1, &stBaseFrmInfo, s32MilliSec);
        SAMPLE_CHECK_EXPR_GOTO(s32Ret != HI_SUCCESS, EXT_RELEASE,
            "Error(%#x),HI_MPI_VPSS_GetChnFrame failed, VPSS_GRP(%d), VPSS_CHN(%d)!\n", s32Ret, s32VpssGrp,
            as32VpssChn[0]);

        /* Do the ai process work from vpss chn1, output to vo venc */
        s32Ret = Yolov5DectionAiProcess(stBaseFrmInfo, stExtFrmInfo);
        SAMPLE_CHECK_EXPR_GOTO(s32Ret != HI_SUCCESS, BASE_RELEASE,
            "Error(%#x),HI_MPI_VPSS_GetChnFrame failed, VPSS_GRP(%d), VPSS_CHN(%d)!\n", s32Ret, s32VpssGrp,
            as32VpssChn[0]);
        
        /* Push VPSS to VENC */
        if (s_stRfcnSwitch.bVenc == HI_TRUE) {
            s32Ret = HI_MPI_VENC_SendFrame(0, &stExtFrmInfo, 0);
            SAMPLE_CHECK_EXPR_GOTO(s32Ret != HI_SUCCESS, EXT_RELEASE, "HI_MPI_VO_SendFrame failed, Error(%#x)!\n", s32Ret);
        }

        /* Push VPSS to VO */
        if (s_stRfcnSwitch.bVo == HI_TRUE) {
            s32Ret = HI_MPI_VO_SendFrame(0, 0, &stExtFrmInfo, 0);
            SAMPLE_CHECK_EXPR_GOTO(s32Ret != HI_SUCCESS, EXT_RELEASE, "HI_MPI_VO_SendFrame failed, Error(%#x)!\n", s32Ret);
        }

    EXT_RELEASE:
        s32Ret = HI_MPI_VPSS_ReleaseChnFrame(0, 2, &stExtFrmInfo);
        if (s32Ret != HI_SUCCESS) {
            SAMPLE_PRT("Error(%#x),HI_MPI_VPSS_ReleaseChnFrame failed,Grp(%d) chn(%d)!\n", s32Ret, s32VpssGrp,
                as32VpssChn[1]);
        }

    BASE_RELEASE:
        s32Ret = HI_MPI_VPSS_ReleaseChnFrame(0, 1, &stBaseFrmInfo);
        if (s32Ret != HI_SUCCESS) {
            SAMPLE_PRT("Error(%#x),HI_MPI_VPSS_ReleaseChnFrame failed,Grp(%d) chn(%d)!\n", s32Ret, s32VpssGrp,
                as32VpssChn[0]);
        }

    }

    return HI_NULL;
}

/* function : uart process */
static HI_VOID* HisignallingMsgHandle(HI_VOID* pArgs)
{
    hi_unused(pArgs);

    /* step 1: init uart and open it */
    uartFd = UartOpenInit();
    if (uartFd < 0) {
        MSG_DEBUG(DEBUG_LOG,"[Error] : [HisignallingMsgHandle] uart1 open failed.\n");
    } else {
        MSG_DEBUG(DEBUG_LOG,"[INFO] : [HisignallingMsgHandle] open uart1 successed.\n");
    }
    HI_UCHAR UartRecvBuf[HISIGNALLING_MSG_ONE_FRAME_LEN] = { 0 };
    while (HI_FALSE == s_bNnieStopSignal) {
        if (HisignallingMsgReceive(uartFd, UartRecvBuf, HISIGNALLING_MSG_ONE_FRAME_LEN) != HISIGNALLING_RET_VAL_ERROR) {
            /* Parse recvbuf process */
            SAMPLE_IVE_MUTEX_LOCK(s_stIveKcfInfo.ModifyMutex);
            ParseUartMsgProcess(UartRecvBuf, &s_stIveKcfInfo);
            SAMPLE_IVE_MUTEX_UNLOCK(s_stIveKcfInfo.ModifyMutex);

        }
        /* Sleep 100ms */
        usleep(HISGNALLING_FREE_TASK_TIME * 1000UL);
    }

    return HI_NULL;
}

/* function : Start vi to vo venc */
static HI_VOID* SAMPLE_VO_VENC_H264( HI_VOID* pArgs )
{
    hi_unused(pArgs);
    /* Run the total process, VI -> VPSS -> CH0 -> DRAW RECT -> CH1 -> NNIE -> VPSS_CH0 FOR OUT (RTSP AND VO) */
    HI_S32 fd = 0;
    HI_S32 s32Ret = HI_SUCCESS;
    PIC_SIZE_E enSize[2] = {PIC_1080P, PIC_1080P};  // output size PIC_1080P for rtsp PIC_D1_MIPI

    /* step 1: start vi vpss vo venc for rtsp */
    s_stRfcnSwitch.bVo = HI_TRUE;
    s_stRfcnSwitch.bVenc = HI_TRUE;
    s_stRfcnSwitch.bMipi = HI_FALSE;

    if (s_stRfcnSwitch.bMipi == HI_TRUE) {
        enSize[1] = PIC_D1_MIPI;
        s_stRfcnSwitch.bVo = HI_FALSE;
    }

    s32Ret = SAMPLE_COMM_IVE_StartViVpssVencVo(&s_stViConfig, &s_stRfcnSwitch, enSize, &fd);
    SAMPLE_CHECK_EXPR_GOTO(s32Ret != HI_SUCCESS, END_RFCN_0, "Error(%#x),SAMPLE_COMM_IVE_StartViVpssVencVo failed!\n",
        s32Ret);

    /* step 2: init necessary sturct params */
    (void)SAMPLE_IVE_KcfInit(&s_stIveKcfInfo);

    /* step 3: Create work thread, push frame to vo and venc */
    s_bNnieStopSignal = HI_FALSE;
    s_bRtmpStopSignal = HI_FALSE;
    
    prctl(PR_SET_NAME, "AudioAencProcess", 0, 0, 0);
    pthread_create(&s_hAudioThread, 0, SAMPLE_COMM_IVE_StartAiAencAdecAo, NULL);

    prctl(PR_SET_NAME, "AiProcess", 0, 0, 0);
    pthread_create(&s_hNnieThread, 0, GetVpssChnFrameAiProcess, NULL);

    prctl(PR_SET_NAME, "LbphTrackProcess", 0, 0, 0);
    pthread_create(&s_hLbphThread, 0, IveLbphProcess, NULL);

    prctl(PR_SET_NAME, "UartProcess", 0, 0, 0);
    pthread_create(&s_hUartThread, 0, HisignallingMsgHandle, NULL);

    /* step 4: end thread and release vi vpss vo venc and frame */
    Pause();

    s_bNnieStopSignal = HI_TRUE;
    s_bRtmpStopSignal = HI_TRUE;
    SVP_NNIE_Yolov5_HandleSig();
    SAMPLE_AUDIO_AiAenc_HandleSig();
    SAMPLE_IVE_LBPH_HandleSig();

    /* Stop venc stream */
    if (s_stRfcnSwitch.bVenc == HI_TRUE) {
        SAMPLE_COMM_VENC_StopGetStream();
    }

    /* Stop mipi output */
    if (s_stRfcnSwitch.bMipi == HI_TRUE) {
        SAMPLE_VO_DISABLE_MIPITx(fd);
        SampleCloseMipiTxFd(fd);
        system("echo 0 > /sys/class/gpio/gpio55/value");
    }

    /* Stop work thread */
    pthread_join(s_hAudioThread, NULL);
    s_hAudioThread = 0;

    pthread_join(s_hNnieThread, NULL);
    s_hNnieThread = 0;

    pthread_join(s_hLbphThread, NULL);
    s_hLbphThread = 0;

    pthread_join(s_hUartThread, NULL);
    s_hUartThread = 0;

    SAMPLE_IVE_KcfDeInit(&s_stIveKcfInfo);
    SAMPLE_IVE_LBPH_Stop();

    /* Unload model */
    PauseDoUnloadModel();
    close(uartFd);

END_RFCN_0:
    SAMPLE_COMM_IVE_StopViVpssVencVo(&s_stViConfig, &s_stRfcnSwitch);
END_RFCN:
    return HI_NULL;
}

/* function : Rfcn Vi->VO */
void SampleYolov5ViVpssVencVo( HI_VOID )
{
    HI_U32 s32Ret = HI_SUCCESS;
    int s32MainFd,temp;
    struct timespec ts = {0,33333};
    pthread_t id;

    /* Init ffmpeg setting */
    s32Ret = initFFmpegNetwork();
    if (s32Ret < 0) return;
    printf("[INFO] : [SampleYolov5ViVpssVencVo] initFFmpegNetwork success\n");
    init_h264_sps_pps_header(1 * 1024);
    s_bRtmpStopSignal = HI_FALSE;

    prctl(PR_SET_NAME, "SAMPLE_VO_VENC_H264", 0, 0, 0);
    pthread_create(&id, 0, SAMPLE_VO_VENC_H264, NULL);
    pthread_detach(id);

    while (HI_FALSE == s_bRtmpStopSignal)
    {
        nanosleep(&ts, NULL);
    }

    sleep(2);
    freeH264SpsPpsHeader();
    cleanFFmpegSetting();
    MSG_DEBUG(DEBUG_FFMPEG, "[INFO] : [SampleTestFFmpeg] free frame queue and memeory.\n");
    return;
}

void Application_HandleSig( HI_VOID )
{
    s_bNnieStopSignal = HI_TRUE;
    s_bRtmpStopSignal = HI_TRUE;
}